#include "znx_tcp_proc.h"
#include "znx_worker_proc.h"
#include "comm/znx_array.h"

#define MAX_DO_ACCEPT_ONCE              16
#define ZNX_TCP_ACCEPTER_WATCH_TIMEOUT  1000


char                            *znx_tcp_accept_thread_tag = "tcp_accepter_thread_tag";


static znx_tcp_accepter_t       g_tcp_accepter = {
    0, NULL, NULL
};


static znx_tcp_connect_opt_t default_tcp_conn_opt = {
    .nodelay = 1,
    .quickack = 1,
};


static znx_bool_t
znx_tcp_accept_sockfd_set(int accept_fd, znx_listener_t *l)
{
    int                     opt_val;
    znx_result_t            ret;

    if (l->nodelay) {
        ret = znx_syscall_setsockopt(accept_fd,
            IPPROTO_TCP, TCP_NODELAY, &opt_val, sizeof(int));
        if (ret < 0) {
            ZNX_DEF_LOG_ERROR(NULL, "setsockopt(TCP_NODELAY) %d failed, errno: %d",
                accept_fd, -ret);
            return ZNX_FALSE;
        }
    }

    if (l->quickack) {
        ret = znx_syscall_setsockopt(accept_fd,
            IPPROTO_TCP, TCP_QUICKACK, &opt_val, sizeof(int));
        if (ret < 0) {
            ZNX_DEF_LOG_ERROR(NULL, "setsockopt(TCP_QUICKACK) %d failed, errno: %d",
                accept_fd, -ret);
            return ZNX_FALSE;
        }
    }

    return ZNX_TRUE;
}


// run in worker thread.
static void
znx_tcp_accept_dispatch_handler(znx_dispatch_event_t *dispatch_event)
{
    znx_listener_t          *l;

    l = dispatch_event->save_data_1;
    znx_conn_t *conn = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn, ZNX_CONN_EPOLL_ET);
    conn->fd = (int)(intptr_t)dispatch_event->data;
    conn->l = l;
    conn->read.handler = l->accept_handler;
    conn->read.ready = 1;
    conn->write.ready = 1;
    znx_atomic_int64_add_fetch(&server_conns, 1);
    conn->server = 1;

    // see znx_tcp_accept_dispatch_event_alloc.
    free(dispatch_event);

    // notify using read event.
    conn->read.handler(&conn->read);
}


// run in accept thread.
static znx_dispatch_event_t *
znx_tcp_accept_dispatch_event_alloc(int sockfd, znx_listener_t *l)
{
    znx_dispatch_event_t    *dispatch_event;
    // the znx_thread_malloc interface is not used here
    // because there is no way to reuse memory.
    // it is inefficient to apply for memory from the global every time.
    dispatch_event = malloc(sizeof(znx_dispatch_event_t));
    dispatch_event->data = (void *)(intptr_t)sockfd;
    znx_listener_conn_hold(l);
    dispatch_event->save_data_1 = l;
    dispatch_event->handler = znx_tcp_accept_dispatch_handler;

    return dispatch_event;
}


static void
znx_tcp_accept_handle(znx_event_t *event)
{
    znx_conn_t                  *c;
    znx_listener_t              *l;
    znx_result_t                ret;
    int                         accept_fd;
    znx_dispatch_event_t        *dispatch_event;

    c = event->data;
    l = c->data;

    for (int i = 0; i < MAX_DO_ACCEPT_ONCE; i++) {
        ret = znx_syscall_accept4(c->fd, NULL, NULL, SOCK_NONBLOCK|SOCK_CLOEXEC);
        if (ret < 0) {
            // EAGAIN or other error.
            return;
        }

        accept_fd = (int)ret;

        if (!znx_tcp_accept_sockfd_set(accept_fd, l)) {
            ZNX_DEF_LOG_ERROR(NULL, "set accept fd %d tcp options failed", accept_fd);
            znx_syscall_close(accept_fd);
            continue;
        }

        dispatch_event = znx_tcp_accept_dispatch_event_alloc(accept_fd, l);
        znx_worker_processor_post_event(dispatch_event);
    }
}


static void
znx_tcp_accepter_listeners_update(znx_tcp_loop_data_t *t_loop_data)
{
    if (t_loop_data->version == g_listeners->version) {
        return;
    }

    znx_array_t                 *tmp_arr;
    znx_tcp_listener_item_t     *item, *items, *itemp;
    size_t                      i;
    znx_bool_t                  found;
    znx_queue_t                 *q;
    znx_listener_t              *l;

    tmp_arr = znx_array_create(0, sizeof(znx_tcp_listener_item_t));

    items = t_loop_data->ls->elts;
    for (i = 0; i < t_loop_data->ls->nelts; i++) {
        item = items + i;
        found = ZNX_FALSE;
        q = znx_queue_head(&g_listeners->head);
        while (q != &g_listeners->head) {
            l = znx_queue_data(q, znx_listener_t, queue);
            q = q->next;
            if (item->l == l) {
                found = ZNX_TRUE;
                itemp = znx_array_push(tmp_arr);
                *itemp = *item;
                break;
            }
        }

        if (found) {
            continue;
        }

        znx_event_loop_remove_conn(&item->conn);
        znx_conn_close(&item->conn);
        znx_listener_free(item->l);
    }

    q = znx_queue_head(&g_listeners->head);
    while (q != &g_listeners->head) {
        l = znx_queue_data(q, znx_listener_t, queue);
        q = q->next;

        if (l->type != SOCK_STREAM) {
            continue;
        }

        found = ZNX_FALSE;
        items = t_loop_data->ls->elts;
        for (i = 0; i < t_loop_data->ls->nelts; i++) {
            item = items + i;
            while (l == item->l) {
                found = ZNX_TRUE;
                break;
            }
        }

        if (found) {
            continue;
        }

        itemp = znx_array_push(tmp_arr);
        znx_conn_init(&itemp->conn, ZNX_CONN_EPOLL_LT);
        itemp->conn.fd = l->sockfd;
        itemp->conn.shared = 1;
        itemp->conn.data = l;
        itemp->conn.read.handler = znx_tcp_accept_handle;
        znx_listener_hold(l);
        itemp->l = l;

        // TODO: we have no way to handle the exception here.
    
        if (!znx_event_loop_start_io(&itemp->conn)) {
            ZNX_DEF_LOG_ERROR(NULL, "start tcp listener %s io failed", l->listen_addr);
            continue;
        }

        if (!znx_event_loop_enable_read(&itemp->conn)) {
            ZNX_DEF_LOG_ERROR(NULL, "enabled tcp listener %s io failed", l->listen_addr);
            continue;
        }
    }

    znx_array_destroy(t_loop_data->ls);
    t_loop_data->ls = tmp_arr;
    t_loop_data->version = g_listeners->version;
}


static void
znx_tcp_accepter_listeners_watch_handle(znx_event_t *event)
{
    znx_tcp_loop_data_t *t_loop_data;
    t_loop_data = event->data;

    znx_rwlock_rlock(&g_listeners->rwlock);
    znx_tcp_accepter_listeners_update(t_loop_data);
    znx_rwlock_runlock(&g_listeners->rwlock);
}


static znx_bool_t
znx_tcp_accepter_start_pool_listener(znx_tcp_loop_data_t *t_loop_data)
{
    znx_queue_t                     *q;
    znx_listener_t                  *l;
    znx_tcp_listener_item_t         *item; 

    q = znx_queue_head(&g_listeners->head);

    while (q != &g_listeners->head) {
        l = znx_queue_data(q, znx_listener_t, queue);
        q = q->next;
        if (l->type != SOCK_STREAM) {
            continue;
        }

        item = znx_array_push(t_loop_data->ls);
        znx_conn_init(&item->conn, ZNX_CONN_EPOLL_LT);
        item->conn.fd = l->sockfd;
        item->conn.shared = 1;
        item->conn.data = l;
        item->conn.read.handler = znx_tcp_accept_handle;
        znx_listener_hold(l);
        item->l = l;

        if (!znx_event_loop_start_io(&item->conn)) {
            ZNX_DEF_LOG_ERROR(NULL, "start tcp listener %s io failed", l->listen_addr);
            return ZNX_FALSE;
        }

        if (!znx_event_loop_enable_read(&item->conn)) {
            ZNX_DEF_LOG_ERROR(NULL, "enabled tcp listener %s io failed", l->listen_addr);
            return ZNX_FALSE;
        }
    }

    t_loop_data->version = g_listeners->version;

    return ZNX_TRUE;
}


// run in accept thread.
static void
znx_tcp_accepter_clear(void *data)
{
    znx_tcp_loop_data_t             *t_loop_data;
    znx_tcp_listener_item_t         *item, *items;
    size_t                          i;

    t_loop_data = data;

    items = t_loop_data->ls->elts;
    for (i = 0; i < t_loop_data->ls->nelts; i++) {
        item = items + i;
        znx_event_loop_remove_conn(&item->conn);
        znx_conn_close(&item->conn);
        znx_listener_free(item->l);
        item->l = NULL;
    }

    znx_array_destroy(t_loop_data->ls);
    znx_thread_free(t_loop_data);
}


// run in accept thread.
static void
znx_tcp_accepter_loop_init(znx_event_loop_data_t *loop_data)
{
    znx_bool_t                      res;
    znx_tcp_loop_data_t             *t_loop_data;
    t_loop_data = znx_thread_calloc(sizeof(znx_tcp_loop_data_t));
    loop_data->data = t_loop_data;
    t_loop_data->ls = znx_array_create(0, sizeof(znx_tcp_listener_item_t));
    znx_rwlock_rlock(&g_listeners->rwlock);
    res = znx_tcp_accepter_start_pool_listener(t_loop_data);
    znx_rwlock_runlock(&g_listeners->rwlock);

    if (!res) {
        ZNX_DEF_LOG_FATAL(NULL, "start watch listener failed");
    }

    // start watch timer task.
    t_loop_data->event.data = t_loop_data;
    t_loop_data->event.handler = znx_tcp_accepter_listeners_watch_handle;
    znx_event_loop_add_loop_once_event(&t_loop_data->event);

    znx_event_loop_cleanup_add(znx_tcp_accepter_clear, t_loop_data);
}


void
znx_tcp_accepter_start(int nthreads)
{
    if (nthreads < 0) {
        ZNX_DEF_LOG_FATAL(NULL, "tcp accept threads must be greate than 0");
    }

    znx_event_loop_data_t       *datas;
    znx_thread_base_t           *threads;

    datas = znx_thread_calloc(
        sizeof(znx_event_loop_data_t) * (size_t)nthreads);
    threads = znx_thread_calloc(
        sizeof(znx_thread_base_t) * (size_t)nthreads);

    char thread_name[64] = { 0 };

    for (int i = 0; i < nthreads; i++) {
        snprintf(thread_name, 64, "tcp-accepter-%d", i);

        znx_event_loop_data_init(
            datas+i,
            znx_tcp_accepter_loop_init,
            NULL,
            NULL);

        znx_thread_base_start(
            znx_tcp_accept_thread_tag,
            threads+i,
            znx_event_loop_runner,
            datas+i,
            0, 0);

        znx_thread_set_name((threads+i)->tid, thread_name);
    }

    g_tcp_accepter.nthreads = nthreads;
    g_tcp_accepter.threads = threads;
    g_tcp_accepter.datas = datas;
}


// run in main thread.
void
znx_tcp_accepter_stop()
{
    if (g_tcp_accepter.nthreads == 0) {
        return;
    }

    int                 i;
    for (i = 0; i < g_tcp_accepter.nthreads; i++) {
        znx_event_loop_quit(g_tcp_accepter.threads+i);
        znx_event_loop_stop(g_tcp_accepter.threads+i);
    }

    for (i = 0; i < g_tcp_accepter.nthreads; i++) {
        znx_thread_base_join(g_tcp_accepter.threads+i, NULL);
    }

    znx_thread_free(g_tcp_accepter.datas);
    znx_thread_free(g_tcp_accepter.threads);
}


void
znx_tcp_accepter_notify()
{
    if (g_tcp_accepter.nthreads == 0) {
        return;
    }

    int                 i;
    for (i = 0; i < g_tcp_accepter.nthreads; i++) {
        znx_event_loop_notify(g_tcp_accepter.datas + i);
    }
}


////////////////// tcp connect ////////////////

typedef struct {
    znx_tcp_connect_handler_pt          connect_handler;
    void                                *user_data;
    void                                *conn_data;
} znx_tcp_connect_data_t;


static void
znx_tcp_connect_done(znx_conn_t *conn, znx_result_t res)
{
    znx_tcp_connect_data_t      *data;
    znx_tcp_connect_handler_pt  connect_handler;
    void                        *user_data;

    data = conn->data;
    conn->data = data->conn_data;
    connect_handler = data->connect_handler;
    user_data = data->user_data;
    znx_thread_free(data);

    if (res == ZNX_SUCCESS) {
        conn->read.ready = 1;
        conn->write.ready = 1;
        znx_event_loop_deactive_conn(conn);
        znx_atomic_int64_add_fetch(&client_conns, 1);
        conn->client = 1;
    } else {
        if (conn->fd != INVALID_FD) {
            znx_event_loop_remove_conn(conn);
            znx_syscall_close(conn->fd);
            conn->fd = INVALID_FD;
        }
    }

    connect_handler(user_data, conn, res);
}


static void
znx_tcp_conn_wait_connect_handler(znx_event_t *event)
{
    znx_conn_t                          *conn;
    znx_result_t                        ret;

    conn = event->data;

    if (event->timedout) {
        ZNX_DEF_LOG_DEBUG(NULL, "connect timedout, fd: %d", conn->fd);
        znx_tcp_connect_done(conn, ZNX_TIMEOUT);
        return;
    }

    int opt_val     = 0;
    int opt_size    = sizeof(opt_val);
    ret = znx_syscall_getsockopt(conn->fd, SOL_SOCKET, SO_ERROR, &opt_val, (socklen_t *)&opt_size);
    if (ret < 0 || opt_val != 0) {
        ZNX_DEF_LOG_DEBUG(NULL, "connect failed, fd: %d, errno: %d, opt_val: %d",
            conn->fd, -ret, opt_val);
        znx_tcp_connect_done(conn, ZNX_FAILED);
        return;
    }

    znx_tcp_connect_done(conn, ZNX_SUCCESS);
}


void
znx_tcp_connect_setup(znx_conn_t *conn,
    znx_address_t *address, znx_address_t *bind_address,
    znx_tcp_connect_opt_t *opts, int64_t timeout_ms,
    znx_tcp_connect_handler_pt connect_handler, void *user_data)
{
    assert(current_thread_is_worker_thread());

    if (opts == NULL) {
        opts = &default_tcp_conn_opt;
    }

    int                     fd;
    znx_result_t            ret;
    znx_tcp_connect_data_t  *data;

    data = znx_thread_calloc(sizeof(znx_tcp_connect_data_t));
    data->connect_handler = connect_handler;
    data->user_data = user_data;
    data->conn_data = conn->data;
    conn->data = data;

    ret = znx_syscall_socket(address->sockaddr.sa_family, SOCK_STREAM, 0);
    if (ret < 0) {
        ZNX_DEF_LOG_ERROR(NULL, "socket() failed, errno: %d", -ret);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    fd = (int)ret;
    conn->fd = fd;

    int opt_val = 1;
    if (opts->nodelay) {
        ret = znx_syscall_setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &opt_val, sizeof(int));
        if (ret < 0) {
            ZNX_DEF_LOG_ERROR(NULL, "setsockopt(TCP_NODELAY) %d failed, errno: %d", fd, -ret);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    if (opts->quickack) {
        ret = znx_syscall_setsockopt(fd, IPPROTO_TCP, TCP_QUICKACK, &opt_val, sizeof(int));
        if (ret < 0) {
            ZNX_DEF_LOG_ERROR(NULL, "setsockopt(TCP_QUICKACK) %d failed, errno: %d", fd, -ret);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    ret = znx_syscall_set_nonblock(fd);
    if (ret < 0) {
        ZNX_DEF_LOG_ERROR(NULL, "setnonblock %d failed, errno: %d", fd, -ret);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    ret = znx_syscall_set_close_exec(fd);
    if (ret < 0) {
        ZNX_DEF_LOG_ERROR(NULL, "setcloseexec %d failed, errno: %d", fd, ret);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    if (bind_address != NULL) {
        ret = znx_syscall_bind(fd, &bind_address->sockaddr,
            (socklen_t)znx_address_socklen(bind_address));
        if (ret < 0) {
            ZNX_DEF_LOG_ERROR(NULL, "bind %d local address failed, errno: %d", fd, -ret);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    ret = znx_syscall_connect(fd, &address->sockaddr, (socklen_t)znx_address_socklen(address));
    // connect success
    if (ret == 0) {
        return znx_tcp_connect_done(conn, ZNX_SUCCESS);
    }

    if (-ret != EINPROGRESS) {
        ZNX_DEF_LOG_ERROR(NULL, "connect() failed, fd: %d, errno: 5d", fd, -ret);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    conn->write.handler = znx_tcp_conn_wait_connect_handler;
    if (timeout_ms > 0) {
        znx_event_loop_add_timer(&conn->write, timeout_ms);
    }

    // wait event loop callback.
    znx_event_loop_start_io(conn);
    znx_event_loop_enable_write(conn);
}


void
znx_tcp_connect_cancel(znx_conn_t *conn)
{
    znx_tcp_connect_data_t          *data;

    if (conn->fd == INVALID_FD) {
        return;
    }

    data = conn->data;
    conn->data = data->conn_data;
    znx_thread_free(data);

    znx_event_loop_remove_conn(conn);
    znx_syscall_close(conn->fd);
    conn->fd = INVALID_FD;
}

